import torch
from torch import nn
from d2l import torch as d2l
import math

def predict(prefix,num_preds,net,vocab,device):
    state = net.begin_state(batch_size=1,device=device)
    outputs = [vocab[prefix[0]]]
    get_input = lambda:torch.tensor([outputs[-1]],device=device).reshape((1,1))
    for y in prefix[1:]:
        _,state = net(get_input(),state)
        outputs.append(vocab[y])
    for _ in range(num_preds):
        y,state = net(get_input(),state)
        outputs.append(int(y.argmax(dim=1).reshape(1)))
    return ''.join([vocab.idx2token[i] for i in outputs])

def grad_clipping(net,theta):
    if isinstance(net,nn.Module):
        params = [p for p in net.parameters() if p.requires_grad]
    else:
        params = net.params
    norm = torch.sqrt(sum(torch.sum((p.grad**2)) for p in params))
    if norm > theta:
        for p in params:
            p.grad[:] *= theta / norm

def train_epoch(net,train_iter,loss,updater,device,use_random_iter=True):
    state, timer = None,d2l.Timer()
    metric = d2l.Accumulator(2)
    for X,Y in train_iter:
        if state is None or use_random_iter:
            state = net.begin_state(batch_size=X.shape[0],device=device)
        else:
            if isinstance(net,nn.Module) and not isinstance(state,tuple):
                state.detach_()
            else:
                for s in state:
                    s.detach_()
        y = Y.T.reshape(-1)
        X = X.to(device)
        y = y.to(device)
        y_hat, state = net(X,state)
        l = loss(y_hat, y.long()).mean()
        if isinstance(updater,torch.optim.Optimizer):
            updater.zero_grad()
            l.backward()
            grad_clipping(net,1)
            updater.step()
        else:
            l.backward()
            grad_clipping(net,1)
            updater(batch_size=1)
        metric.add(l * y.numel(), y.numel())
    return math.exp(metric[0] / metric[1]), metric[1] / timer.stop()

def train(net, train_iter, vocab, lr, num_epochs, device, use_random_iter=True):
    loss = nn.CrossEntropyLoss()

    animator = d2l.Animator(xlabel='epoch',ylabel='perplexity',
                            legend=['train'],xlim=[10,num_epochs])

    if isinstance(net,nn.Module):
        updater = torch.optim.SGD(net.parameters(),lr)
    else:
        updater = lambda batch_size: d2l.sgd(net.params,lr,batch_size)
    pred = lambda prefix:predict(prefix,50,net,vocab,device)
    for epoch in range(num_epochs):
        ppl, speed = train_epoch(net,train_iter,loss,updater,device,use_random_iter)
        if (epoch + 1) % 10 == 0:
            print(pred('time traveller'))
            animator.add(epoch + 1, [ppl])
    print('困惑度:',ppl,speed, '标记/秒')
    print(pred('time traveller'))
    print(pred('traveller'))
